package com.by.toolkit.utils;

import com.google.common.collect.Maps;
import org.slf4j.MDC;

import java.util.Map;
import java.util.Objects;

/**
 * 线程上下文工具类
 *
 * @author YanBingHao
 * @since 2024/1/4
 */
public class ThreadLocalUtil {

    private static final ThreadLocal<Map<String, Object>> THREAD_LOCAL = new InheritableThreadLocal<Map<String, Object>>(){
        @Override
        protected Map<String, Object> initialValue() {
            return Maps.newHashMap();
        }
    };

    static {
        THREAD_LOCAL.set(Maps.newHashMap());
    }

    private ThreadLocalUtil() {
    }

    private static class ThreadLocalUtilHolder {
        public static final ThreadLocalUtil INSTANCE = new ThreadLocalUtil();
    }

    public static void put(String key, Object val, boolean mdc) {
        Map<String, Object> context = THREAD_LOCAL.get();
        if (Objects.nonNull(context) && Objects.nonNull(val)) {
            context.put(key, val);
            if (mdc) {
                MDC.put(key, val.toString());
            }
        }
    }

    public static String getString(String key) {
        Map<String, Object> context = THREAD_LOCAL.get();
        if (Objects.isNull(context)) {
            return null;
        }
        Object o = context.get(key);
        if (Objects.isNull(o)) {
            return null;
        }
        return o.toString();
    }

    public static Long getLong(String key) {
        Map<String, Object> context = THREAD_LOCAL.get();
        if (Objects.isNull(context)) {
            return null;
        }
        Object o = context.get(key);
        if (Objects.isNull(o)) {
            return null;
        }
        return Long.parseLong(o.toString());
    }

    public static Integer getInt(String key) {
        Map<String, Object> context = THREAD_LOCAL.get();
        if (Objects.isNull(context)) {
            return null;
        }
        Object o = context.get(key);
        if (Objects.isNull(o)) {
            return null;
        }
        if (o instanceof Integer) {
            return (Integer) o;
        }
        return Integer.parseInt(o.toString());
    }

    public static void removeKey(String key, boolean mdc) {
        Map<String, Object> context = THREAD_LOCAL.get();
        if (Objects.isNull(context)) {
            return;
        }
        context.remove(key);
        if (mdc) {
            MDC.remove(key);
        }
    }

    public static void remove(boolean mdc) {
        THREAD_LOCAL.remove();
    }

}
